# 池化层
import numpy as np
from Layer import Layer


class MeanPool(Layer):
    def __init__(self, input_shape, pooling_size, pooling_stride):
        super().__init__()
        self.pooling_size = pooling_size  # 池化操作的采样区域边长
        self.stride = pooling_stride  # 池化操作移动的步长
        self.max_index = []
        self.input_shape = input_shape
        self.output_shape = []
        self.pad_shape = []
        self.cal_output_shape()

    # 如果模型没有经过整体考虑，比如采样窗口和步长是随意设置的，那边在池化前还是需要进行一次padding的----尽量避免
    def padding4Pooling(self, oriImage):
        oriHeight, oriWidth = oriImage.shape
        newHeight = np.ceil(oriHeight / self.stride) * self.stride
        newWidth = np.ceil(oriWidth / self.stride) * self.stride
        newHeight, newWidth = int(newHeight), int(newWidth)
        newImage = np.zeros((newHeight, newWidth))  # 创建一个0矩阵，尺寸与填充后的图像大小相同
        newImage[0: oriHeight, 0: oriWidth] = oriImage  # 把原始图像覆盖到0矩阵的中心位置
        return newImage

    # 基于出入图像的尺寸和数量，以及卷积核等的情况，计算输出的图像的个数和尺寸
    def cal_output_shape(self):
        input_image = np.zeros(self.input_shape)
        output_image = self.forward(input_image)
        self.output_shape = output_image.shape

    # 前向传播，即执行池化,使用的是平局池化(mean pooling)
    def forward(self, input_image):
        self.input_shape = input_image.shape
        pad_image = self.padding4Pooling((input_image[0][0]))
        new_height, new_width = pad_image.shape
        self.pad_shape = [input_image.shape[0], input_image.shape[1], new_height, new_width]
        output_kchw = []
        for b in range(0, input_image.shape[0]):
            output_chw = []
            for c in range(0, input_image.shape[1]):
                pad_image = self.padding4Pooling((input_image[b][c]))
                output_hw = []
                for h in range(0, new_height, self.stride):
                    output_w = []
                    for w in range(0, new_width, self.stride):
                        slice_of_image = pad_image[h: h + self.stride, w: w + self.stride]
                        mean_v = np.mean(slice_of_image)
                        output_w.append(mean_v)
                    output_hw.append(output_w)
                output_hw = np.array(output_hw)
                output_chw.append(output_hw)
            output_chw = np.array(output_chw)
            output_kchw.append(output_chw)
        output_kchw = np.array(output_kchw)
        self.output_shape = output_kchw.shape
        return output_kchw

    # 反向传播，基于池化层输出误差，反推计算卷积核输出像素点的误差,使用的是平局池化(mean pooling)
    # 对于平局池化，如果由2x2，池化为1x1，那么反向传播时，仅需将该1x1的误差值，平均填充到2x2的区域上
    # 例如 池化后是值是8，那么池化前2x2的区域的值为[[2 2],[2 2]]
    def backward(self, errorFromLaterLayer):
        back_output = np.zeros(self.pad_shape)
        for b in range(back_output.shape[0]):
            for c in range(back_output.shape[1]):
                for h in range(0, back_output.shape[2], self.stride):  # 遍历一副图像的所有像素点(不包括填充部分)
                    for w in range(0, back_output.shape[3], self.stride):
                        back_output[b, c, h: h + self.pooling_size, w: w + self.pooling_size] = \
                            np.ones((self.pooling_size, self.pooling_size)) * (errorFromLaterLayer[
                                    b, c, int(h / self.stride), int(w / self.stride)] / (
                                    self.pooling_size * self.pooling_size))
        back_output = back_output[:, :, 0:self.input_shape[2], 0:self.input_shape[3]]
        return back_output


class MaxPool(Layer):
    def __init__(self, input_shape, pooling_size, pooling_stride):
        super().__init__()
        self.stride = (pooling_stride, pooling_stride)
        self.pooling_shape = (pooling_size, pooling_size)
        self.max_index = []
        self.input_shape = input_shape
        self.output_shape = []
        self.cal_output_shape()

    # 基于出入图像的尺寸和数量，以及卷积核等的情况，计算输出的图像的个数和尺寸
    def cal_output_shape(self):
        input_image = np.zeros(self.input_shape)
        output_image = self.forward(input_image)
        self.output_shape = output_image.shape

    def forward(self, input_tensor):
        self.input_shape = input_tensor.shape
        self.output_shape = (input_tensor.shape[0], input_tensor.shape[1],
                             (input_tensor.shape[2] - self.pooling_shape[0]) // self.stride[0] + 1,
                             (input_tensor.shape[3] - self.pooling_shape[1]) // self.stride[1] + 1)

        output = np.empty(self.output_shape)
        self.max_index = np.zeros(self.output_shape, dtype=int)

        for b in range(self.output_shape[0]):
            for c in range(self.output_shape[1]):
                for h in range(self.output_shape[2]):
                    for w in range(self.output_shape[3]):
                        pooling_x = h * self.stride[0]
                        pooling_y = w * self.stride[1]
                        pooling_field = input_tensor[b, c, pooling_x:pooling_x + self.pooling_shape[0],
                                        pooling_y:pooling_y + self.pooling_shape[1]]
                        output[b, c, h, w] = np.max(pooling_field)
                        self.max_index[b, c, h, w] = np.argmax(pooling_field)

        return output

    def backward(self, error_tensor):
        error_extend = np.zeros(self.input_shape)

        for b in range(self.input_shape[0]):
            for c in range(self.input_shape[1]):
                for h in range(self.output_shape[2]):
                    for w in range(self.output_shape[3]):
                        back_x = h * self.stride[0]
                        back_y = w * self.stride[1]
                        pooling_field = error_extend[b, c, back_x:back_x + self.pooling_shape[0],
                                        back_y:back_y + self.pooling_shape[1]]
                        index0 = self.max_index[b, c, h, w] // self.pooling_shape[0]
                        index1 = self.max_index[b, c, h, w] % self.pooling_shape[1]
                        pooling_field[index0, index1] += error_tensor[b, c, h, w]
                        error_extend[b, c, back_x:back_x + self.pooling_shape[0],
                        back_y:back_y + self.pooling_shape[1]] = pooling_field

        error_tensor = error_extend
        return error_tensor
